{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fc48f5ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import argparse\n",
    "import mindspore.dataset as ds\n",
    "import mindspore.nn as nn\n",
    "from mindspore.train import Model\n",
    "from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor\n",
    "from mindspore.train.serialization import load_checkpoint, load_param_into_net\n",
    "from mindspore.common.initializer import Normal\n",
    "from mindspore import context\n",
    "import mindspore.dataset.vision.c_transforms as CV\n",
    "import mindspore.dataset.transforms.c_transforms as C\n",
    "from mindspore.dataset.vision import Inter\n",
    "from mindspore.nn import Accuracy\n",
    "from mindspore.common import dtype as mstype\n",
    "from mindspore.nn import SoftmaxCrossEntropyWithLogits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e125b1be",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset(data_path, batch_size=32, repeat_size=1,\n",
    "                   num_parallel_workers=1):\n",
    "    \"\"\" create dataset for train or test\n",
    "    Args:\n",
    "        data_path: Data path\n",
    "        batch_size: The number of data records in each group\n",
    "        repeat_size: The number of replicated data records\n",
    "        num_parallel_workers: The number of parallel workers\n",
    "    \"\"\"\n",
    "    # define dataset\n",
    "    mnist_ds = ds.MnistDataset(data_path)\n",
    "\n",
    "    # define operation parameters\n",
    "    resize_height, resize_width = 32, 32\n",
    "    rescale = 1.0 / 255.0\n",
    "    shift = 0.0\n",
    "    rescale_nml = 1 / 0.3081\n",
    "    shift_nml = -1 * 0.1307 / 0.3081\n",
    "\n",
    "    # define map operations\n",
    "    resize_op = CV.Resize((resize_height, resize_width), interpolation=Inter.LINEAR)  # Resize images to (32, 32)\n",
    "    rescale_nml_op = CV.Rescale(rescale_nml, shift_nml) # normalize images\n",
    "    rescale_op = CV.Rescale(rescale, shift) # rescale images\n",
    "    hwc2chw_op = CV.HWC2CHW() # change shape from (height, width, channel) to (channel, height, width) to fit network.\n",
    "    type_cast_op = C.TypeCast(mstype.int32) # change data type of label to int32 to fit network\n",
    "\n",
    "    # apply map operations on images\n",
    "    mnist_ds = mnist_ds.map(operations=type_cast_op, input_columns=\"label\", num_parallel_workers=num_parallel_workers)\n",
    "    mnist_ds = mnist_ds.map(operations=resize_op, input_columns=\"image\", num_parallel_workers=num_parallel_workers)\n",
    "    mnist_ds = mnist_ds.map(operations=rescale_op, input_columns=\"image\", num_parallel_workers=num_parallel_workers)\n",
    "    mnist_ds = mnist_ds.map(operations=rescale_nml_op, input_columns=\"image\", num_parallel_workers=num_parallel_workers)\n",
    "    mnist_ds = mnist_ds.map(operations=hwc2chw_op, input_columns=\"image\", num_parallel_workers=num_parallel_workers)\n",
    "\n",
    "    # apply DatasetOps\n",
    "    buffer_size = 10000\n",
    "    mnist_ds = mnist_ds.shuffle(buffer_size=buffer_size)  # 10000 as in LeNet train script\n",
    "    mnist_ds = mnist_ds.batch(batch_size, drop_remainder=True)\n",
    "    mnist_ds = mnist_ds.repeat(repeat_size)\n",
    "\n",
    "    return mnist_ds\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1094d3b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LeNet5(nn.Cell):\n",
    "    \"\"\"Lenet network structure.\"\"\"\n",
    "    # define the operator required\n",
    "    def __init__(self, num_class=10, num_channel=1):\n",
    "        super(LeNet5, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(num_channel, 6, 5, pad_mode='valid')\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')\n",
    "        self.fc1 = nn.Dense(16 * 5 * 5, 120, weight_init=Normal(0.02))\n",
    "        self.fc2 = nn.Dense(120, 84, weight_init=Normal(0.02))\n",
    "        self.fc3 = nn.Dense(84, num_class, weight_init=Normal(0.02))\n",
    "        self.relu = nn.ReLU()\n",
    "        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.flatten = nn.Flatten()\n",
    "\n",
    "    # use the preceding operators to construct networks\n",
    "    def construct(self, x):\n",
    "        x = self.max_pool2d(self.relu(self.conv1(x)))\n",
    "        x = self.max_pool2d(self.relu(self.conv2(x)))\n",
    "        x = self.flatten(x)\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7bd97d34",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_net(network_model, epoch_size, data_path, repeat_size, ckpoint_cb, sink_mode):\n",
    "    \"\"\"Define the training method.\"\"\"\n",
    "    print(\"============== Starting Training ==============\")\n",
    "    # load training dataset\n",
    "    ds_train = create_dataset(os.path.join(data_path, \"train\"), 32, repeat_size)\n",
    "    network_model.train(epoch_size, ds_train, callbacks=[ckpoint_cb, LossMonitor()], dataset_sink_mode=sink_mode)\n",
    "\n",
    "\n",
    "def test_net(network, network_model, data_path):\n",
    "    \"\"\"Define the evaluation method.\"\"\"\n",
    "    print(\"============== Starting Testing ==============\")\n",
    "    # load the saved model for evaluation\n",
    "    param_dict = load_checkpoint(\"checkpoint_lenet-1_1875.ckpt\")\n",
    "    # load parameter to the network\n",
    "    load_param_into_net(network, param_dict)\n",
    "    # load testing dataset\n",
    "    ds_eval = create_dataset(os.path.join(data_path, \"test\"))\n",
    "    acc = network_model.eval(ds_eval, dataset_sink_mode=False)\n",
    "    print(\"============== Accuracy:{} ==============\".format(acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dce7c334",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============== Starting Training ==============\n",
      "epoch: 1 step: 1875, loss is 0.018725372850894928\n",
      "epoch: 2 step: 1875, loss is 0.0090101957321167\n",
      "epoch: 3 step: 1875, loss is 0.022994326427578926\n",
      "epoch: 4 step: 1875, loss is 0.13456447422504425\n",
      "epoch: 5 step: 1875, loss is 0.008279167115688324\n",
      "============== Starting Testing ==============\n",
      "============== Accuracy:{'Accuracy': 0.9678485576923077} ==============\n"
     ]
    }
   ],
   "source": [
    "device_target = \"CPU\" #choices=['Ascend', 'CPU']\n",
    "context.set_context(mode=context.GRAPH_MODE, device_target=device_target)\n",
    "dataset_sink_mode = not device_target == \"CPU\"\n",
    "# download mnist dataset\n",
    "# learning rate setting\n",
    "lr = 0.01\n",
    "momentum = 0.9\n",
    "dataset_size = 1\n",
    "mnist_path = \"./MNIST\"\n",
    "# define the loss function\n",
    "net_loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')\n",
    "train_epoch = 5\n",
    "# create the network\n",
    "net = LeNet5()\n",
    "# define the optimizer\n",
    "net_opt = nn.Momentum(net.trainable_params(), lr, momentum)\n",
    "config_ck = CheckpointConfig(save_checkpoint_steps=1875, keep_checkpoint_max=10)\n",
    "# save the network model and parameters for subsequence fine-tuning\n",
    "ckpoint = ModelCheckpoint(prefix=\"checkpoint_lenet\", config=config_ck)\n",
    "# group layers into an object with training and evaluation features\n",
    "model = Model(net, net_loss, net_opt, metrics={\"Accuracy\": Accuracy()})\n",
    "\n",
    "train_net(model, train_epoch, mnist_path, dataset_size, ckpoint, dataset_sink_mode)\n",
    "test_net(net, model, mnist_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "68dd2ffb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindspore.train.serialization import load_checkpoint, save_checkpoint, export\n",
    "import numpy as np\n",
    "from mindspore import Tensor\n",
    "network = LeNet5()\n",
    "load_checkpoint(\"./checkpoint_lenet-1_1875.ckpt\", network)\n",
    "input_data = np.random.uniform(0.0, 1.0, size = [1, 1, 32, 32]).astype(np.float32)\n",
    "export(network, Tensor(input_data), file_name = './mnist', file_format = 'AIR') "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MindSpore",
   "language": "python",
   "name": "mindspore"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
